//
//  HookCore.cpp
//  DemoV2
//
//  Created by 董一龙 on 2022/8/5.
//

#include "HookCore.hpp"
#include "Logger.hpp"
#include "fishhook.h"
#include <sys/time.h>
#include <sys/types.h>
#include <sys/syscall.h>
#include <unistd.h>
#include <dispatch/dispatch.h>
#include <unordered_map>
//#include "hookObjcMsgSend-arm64.s"
#include <thread>
#include "../util/TimeUtil.hpp"
#include "Filter.hpp"
#if defined(DEBUG)
#include "../test/test.hpp"
#endif

// c和汇编部分
#ifdef __cplusplus
extern "C" {
#endif

__attribute__((__naked__)) void hook_Objc_msgSend(void);

#include "hookObjcMsgSendArm64.inline.h"
    #if defined (__arm64__)
        
    #elif defined (__x86_64__)
    //#error "Unsupport Arch x86"
    #elif defined (__ppc64__)
    //#error "Unsupport Arch powerpc"
    #endif

#ifdef __cplusplus
}
#endif

void (*orgin_objc_msgSend)(void);
void (*orgin_objc_msgSendSuper2)(void);
extern "C" {
void hook_msgSend(void);
void hook_msgSendSuper2(void);
}
extern id (*orig_objc_msgSend)(id, SEL, ...);
extern id (*orig_objc_msgSendSuper)(id, SEL, ...);

thread_local std::stack<uintptr_t> HookCore::t_lrStack;
thread_local std::stack<HookCore::ObjcCallRecord> HookCore::t_callRecordStack;
thread_local std::stack<HookCore::LogState> HookCore::t_logStateStack;
thread_local HookCore::LogState HookCore::t_logState = HookCore::LogState::mask_recursion;
thread_local uint32_t HookCore::t_logStateDepth = 0;
thread_local uint64_t HookCore::t_timeOffset = 0;
bool HookCore::s_isEnableLog = true;
Filter HookCore::filter;

extern "C" {
void hook_objc_msgSend_before(id self, SEL cmd, uintptr_t lr);
}
extern "C" {
uintptr_t hook_objc_msgSend_after(BOOL is_objc_msgSendSuper);
}
extern "C" {
static void (*origin_objc_end_catch)(void);
void hook_objc_end_catch(void);
}
void HookCore::startHook() {
    printf("%s", __FUNCTION__);

#if defined(DEBUG)
    test();
#endif
    
    
    static dispatch_once_t onceToken;
        dispatch_once(&onceToken, ^{
            Logger::setLogFileName(HookCore::config().log_filename);
            HookCore::getTimestamp();
            
            filter.batchRegisterClassPrefixHit({
                { .prefix = "_*", .isWhiteList = false, },
                { .prefix = "OS_*", .isWhiteList = false, },
                { .prefix = "nil*", .isWhiteList = false, }
            });
            filter.registerMethodSigntureHit(
                "-[PersonChannelViewController viewDidLoad]",
                (uint32_t)LogState::match_recursion,
                0);
            HookCore::enableLog();
            HookCore::config().only_main_thread = true;
            
            struct rebinding rebindObjc_msgSend;
            rebindObjc_msgSend.name = "objc_msgSend";
            rebindObjc_msgSend.replacement = (void*)hook_msgSend;
            rebindObjc_msgSend.replaced = (void **)&orgin_objc_msgSend;
            struct rebinding rebindObjc_msgSendSuper2;
            rebindObjc_msgSendSuper2.name = "objc_msgSendSuper2";
            rebindObjc_msgSendSuper2.replacement = (void *)hook_msgSendSuper2;
            rebindObjc_msgSendSuper2.replaced = (void **)&orgin_objc_msgSendSuper2;
            struct rebinding rebindObjc_endCatch;
            rebindObjc_endCatch.name = "objc_end_catch";
            rebindObjc_endCatch.replacement = (void *)hook_objc_end_catch;
            rebindObjc_endCatch.replaced = (void **)&origin_objc_end_catch;
            struct rebinding rebs[3] = {rebindObjc_msgSend, rebindObjc_msgSendSuper2, rebindObjc_endCatch};
            rebind_symbols(rebs, 3);

//            rebind_symbols((struct rebinding[6]){
//                {"objc_msgSend", (void *)hook_Objc_msgSend, (void **)&orig_objc_msgSend},
//            }, 1);
        });
}

void HookCore::stopHook() {
    printf("%s", __FUNCTION__);
    s_isEnableLog = false;
}

void HookCore::enableLog() {
    printf("%s", __FUNCTION__);
    s_isEnableLog = true;
}

void HookCore::disableLog() {
    printf("%s", __FUNCTION__);
    s_isEnableLog = false;
}

void before_objc_msgSend(id self, SEL cmd, uintptr_t lr) {
    HookCore::runBeforeObjcMsgSend(self, object_getClass(self), cmd, lr);
}

uintptr_t after_objc_msgSend(){
    return HookCore::runAfterObjcMsgSend(NO);
}


void hook_objc_msgSend_before(id self, SEL cmd, uintptr_t lr) {
    HookCore::runBeforeObjcMsgSend(self, object_getClass(self), cmd, lr);
}

uintptr_t hook_objc_msgSend_after(BOOL is_objc_msgSendSuper)
{
    return HookCore::runAfterObjcMsgSend(is_objc_msgSendSuper);
}

void hook_objc_end_catch(void) {
    origin_objc_end_catch();
    int *p = (int *)__builtin_return_address(1);
    
    while (!HookCore::t_lrStack.empty()) {
        int *lr = (int *)HookCore::t_lrStack.top();
        printf("cur lr = %p\t", lr);
        printf("expect lr = %p\n", p);
        if (lr == p) {
            break; // TODO: 记录异常捕获日志
        }
        // 记录栈展开日志
        HookCore::t_lrStack.pop();
        HookCore::ObjcCallRecord callRecord = HookCore::t_callRecordStack.top();
        HookCore::t_callRecordStack.pop();
        thread_local uint64_t tid = 0;
        if(tid == 0) {
            pthread_threadid_np(NULL, &tid);
        }
        HookCore::ObjcCallRecordLogInfo callRecordLogInfo = {
            .callRecord = callRecord,
            .tid = tid,
            .unionValue.duration = HookCore::getTimestamp() - callRecord.timestamp,
            .logEvent = HookCore::ObjcCallRecordLogInfo::LogEvent::method_duration,
        };
        HookCore::log(callRecordLogInfo);
    }
}

void HookCore::reset() {
    // TODO: 处理抛出异常被捕获时的路径
//    decltype(HookCore::t_lrStack) emptyLrStack{};
//    HookCore::t_lrStack.swap(emptyLrStack);
//    HookCore::t_logStateDepth = 0;
//    decltype(HookCore::t_logStateStack) emptyStateStack{};
//    HookCore::t_logStateStack.swap(emptyStateStack);
//
}
void HookCore::runBeforeObjcMsgSend(id self, Class cls, SEL sel, uintptr_t lr) {
    
    thread_local bool is_main_thread = pthread_main_np();
    if(HookCore::config().only_main_thread && is_main_thread == false) {
        return;
    }
    
    uint64_t start_time = HookCore::getTimestamp();
    HookCore::t_lrStack.push(lr);
    HookCore::ObjcCallRecord callRecord = {
        .sel = sel,
        .cls = cls,
        .isInstanceMethod = !class_isMetaClass(cls),
        .timestamp = HookCore::getTimestamp(),
    };
    HookCore::t_callRecordStack.push(callRecord);
    
    //    检查是否命中，获取事件ID
    HookCore::LogState logState =
        (HookCore::LogState)filter.testMethodSignture(Filter::getMethodSignture(cls, sel));
    updateLogState(logState);
    
    uint64_t end_time = HookCore::getTimestamp();
    HookCore::t_timeOffset += (end_time - start_time);
}

uintptr_t HookCore::runAfterObjcMsgSend(BOOL is_objc_msgSendSuper) {
    thread_local bool is_main_thread = pthread_main_np();
    if(HookCore::config().only_main_thread && is_main_thread == false) {
        return 0;  // 更改汇编后
    }
    
    uint64_t start_time = HookCore::getTimestamp();
    uintptr_t lr = HookCore::t_lrStack.top();
    HookCore::t_lrStack.pop();
    HookCore::ObjcCallRecord callRecord = HookCore::t_callRecordStack.top();
    HookCore::t_callRecordStack.pop();
    HookCore::LogState logState = t_logStateStack.top();
    t_logStateStack.pop();
    
    // 不关心的方法不记录
    if(logState == LogState::mask_recursion ||
       logState == LogState::unknown) {
        assert(logState != LogState::unknown && "logState错误");
        return lr;
    }
    
    thread_local uint64_t tid = 0;
    if(tid == 0) {
        pthread_threadid_np(NULL, &tid);
    }
    
//    if(is_objc_msgSendSuper) {
////        printf("Super");
//        callRecord.cls = class_getSuperclass(callRecord.cls);
//    }
    HookCore::ObjcCallRecordLogInfo callRecordLogInfo = {
        .callRecord = callRecord,
        .tid = tid,
        .unionValue.duration = getTimestamp() - callRecord.timestamp,
        .logEvent = HookCore::ObjcCallRecordLogInfo::LogEvent::method_duration,
        .isSuper = is_objc_msgSendSuper,
    };
    
    // 用时过少的方法不记录
    if(callRecordLogInfo.logEvent == HookCore::ObjcCallRecordLogInfo::LogEvent::method_duration &&
       callRecordLogInfo.unionValue.duration == 0) {
        return lr;
    }
    
    HookCore::log(callRecordLogInfo);
    uint64_t end_time = HookCore::getTimestamp();
    HookCore::t_timeOffset += (end_time - start_time);
    return lr;
}

void HookCore::log(struct HookCore::ObjcCallRecordLogInfo const& callRecordLogInfo) {
    if(s_isEnableLog == false) {
        return;
    }
    
    char const *classname = class_getName(callRecordLogInfo.callRecord.cls);
    // 属于黑名单前缀则不记录日志
    if(strcmp(classname, "nil") == 0) {
        
    }
    if(filter.testClassPrefix(classname) < 0) {
        return;
    }
    
    static std::unordered_map<ObjcCallRecordLogInfo::LogEvent, std::function<void(HookCore::ObjcCallRecordLogInfo const& callRecordLogInfo, char *buffer, uint32_t size)>> logEventFormatHandlerMap {
        { ObjcCallRecordLogInfo::LogEvent::method_call, HookCore::formatForBeginEvent },
        { ObjcCallRecordLogInfo::LogEvent::method_return, HookCore::formatForEndEvent },
        { ObjcCallRecordLogInfo::LogEvent::method_duration, HookCore::formatForDurationEvent },
    };
    
    thread_local static char msgBuffer[1024 * 4];
    if(logEventFormatHandlerMap.count(callRecordLogInfo.logEvent) == 0) {
        printf("error\n");
        std::exit(0);
    }
    auto formatHandler = logEventFormatHandlerMap[callRecordLogInfo.logEvent];
    formatHandler(callRecordLogInfo, msgBuffer, sizeof(msgBuffer));
    Logger::getInstance().asyncLog(msgBuffer, (uint32_t)strlen(msgBuffer));
}

HookCore::Config& HookCore::config() {
    static HookCore::Config config;
    return config;
}

void HookCore::formatForBeginEvent(HookCore::ObjcCallRecordLogInfo const& callRecordLogInfo, char *buffer, uint32_t size) {
    snprintf(buffer, size,
             R"({"name": "%s", "ph": "%s", "pid": "%s", "tid": "%llu", "ts": %llu, "sf": %lu})" "\n",
             HookCore::getMethodSignture(callRecordLogInfo.callRecord.cls, callRecordLogInfo.callRecord.sel, callRecordLogInfo.isSuper).c_str(),
             "B",
             "Main",
             callRecordLogInfo.tid,
             callRecordLogInfo.callRecord.timestamp,
             t_logStateStack.size() + 1);
}

void HookCore::formatForEndEvent(HookCore::ObjcCallRecordLogInfo const& callRecordLogInfo, char *buffer, uint32_t size) {
    snprintf(buffer, size,
             R"({"name": "%s", "ph": "%s", "pid": "%s", "tid": "%llu", "ts": %llu, "sf": %lu})" "\n",
             HookCore::getMethodSignture(callRecordLogInfo.callRecord.cls, callRecordLogInfo.callRecord.sel, callRecordLogInfo.isSuper).c_str(),
             "E",
             "Main",
             callRecordLogInfo.tid,
             callRecordLogInfo.unionValue.timestamp,
             t_logStateStack.size() + 1);}

void HookCore::formatForDurationEvent(HookCore::ObjcCallRecordLogInfo const& callRecordLogInfo, char *buffer, uint32_t size) {
    snprintf(buffer, size,
             R"({"name": "%s", "ph": "%s", "pid": "%s", "tid": "%llu", "ts": %llu, "dur": %llu, "sf": %lu})" "\n",
             HookCore::getMethodSignture(callRecordLogInfo.callRecord.cls, callRecordLogInfo.callRecord.sel, callRecordLogInfo.isSuper).c_str(),
             "X",
             "Main",
             callRecordLogInfo.tid,
             callRecordLogInfo.callRecord.timestamp,
             callRecordLogInfo.unionValue.duration,
             t_logStateStack.size() + 1);
}

void HookCore::updateLogState(HookCore::LogState logState) {
    if(!t_logStateStack.empty() && t_logStateDepth >= t_logStateStack.size()) {
        t_logState = t_logStateStack.top();
        t_logStateDepth = (uint32_t)t_logStateStack.size();
    }
    
    if(logState == LogState::mask_recursion || logState == LogState::match_recursion) {
        t_logState = logState;
        t_logStateStack.push(logState);
        t_logStateDepth = (uint32_t)t_logStateStack.size();
    }
    else if(logState == LogState::unknown){
        t_logStateStack.push(t_logState);
    }
    else {
        t_logStateStack.push(LogState::unknown);
        assert(0 && "Error LogState\n");
    }
}

/// 获得相对时间，毫秒时间戳
uint64_t HookCore::getTimestamp() {
    static uint64_t const start_timestamp = TimeUtil::getTimeStamp();
    return TimeUtil::getTimeStamp() - start_timestamp - HookCore::t_timeOffset;
}

std::string HookCore::getMethodSignture(Class cls, SEL sel, bool isSuper = false) {
    char const *classname = class_getName(cls);
    char const *methodname = sel_getName(sel);
    bool isInstanceMethod = class_isMetaClass(cls);
    thread_local static char nameBuffer[1024];
    if(isSuper) {
        char const *superClassname = class_getName(class_getSuperclass(cls));
        snprintf(nameBuffer, sizeof(nameBuffer),
                 "%s[(super %s)%s %s]",
                 isInstanceMethod ? "+" : "-",
                 classname,
                 superClassname,
                 methodname);
    }
    else {
        snprintf(nameBuffer, sizeof(nameBuffer),
                 "%s[%s %s]",
                 isInstanceMethod ? "+" : "-",
                 classname,
                 methodname);
    }
    
    return std::string{ nameBuffer };
}
